{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X80i_girFR2o"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "bB8gHCR3FVC0"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z17OmgavQfp4"
      },
      "source": [
        "# Recommending movies: retrieval\n",
        "\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/recommenders/examples/basic_retrieval\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/recommenders/blob/main/docs/examples/basic_retrieval.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/recommenders/blob/main/docs/examples/basic_retrieval.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/recommenders/docs/examples/basic_retrieval.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kCeYA79m1DEX"
      },
      "source": [
        "Real-world recommender systems are often composed of two stages:\n",
        "\n",
        "1. The retrieval stage is responsible for selecting an initial set of hundreds of candidates from all possible candidates. The main objective of this model is to efficiently weed out all candidates that the user is not interested in. Because the retrieval model may be dealing with millions of candidates, it has to be computationally efficient.\n",
        "2. The ranking stage takes the outputs of the retrieval model and fine-tunes them to select the best possible handful of recommendations. Its task is to narrow down the set of items the user may be interested in to a shortlist of likely candidates.\n",
        "\n",
        "In this tutorial, we're going to focus on the first stage, retrieval. If you are interested in the ranking stage, have a look at our [ranking](basic_ranking) tutorial.\n",
        "\n",
        "Retrieval models are often composed of two sub-models:\n",
        "\n",
        "1. A query model computing the query representation (normally a fixed-dimensionality embedding vector) using query features.\n",
        "2. A candidate model computing the candidate representation (an equally-sized vector) using the candidate features\n",
        "\n",
        "The outputs of the two models are then multiplied together to give a query-candidate affinity score, with higher scores expressing a better match between the candidate and the query.\n",
        "\n",
        "In this tutorial, we're going to build and train such a two-tower model using the Movielens dataset.\n",
        "\n",
        "We're going to:\n",
        "\n",
        "1. Get our data and split it into a training and test set.\n",
        "2. Implement a retrieval model.\n",
        "3. Fit and evaluate it.\n",
        "4. Export it for efficient serving by building an approximate nearest neighbours (ANN) index.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H7QZ3kkMQo48"
      },
      "source": [
        "## The dataset\n",
        "\n",
        "The Movielens dataset is a classic dataset from the [GroupLens](https://grouplens.org/datasets/movielens/) research group at the University of Minnesota. It contains a set of ratings given to movies by a set of users, and is a workhorse of recommender system research.\n",
        "\n",
        "The data can be treated in two ways:\n",
        "\n",
        "1. It can be interpreted as expressesing which movies the users watched (and rated), and which they did not. This is a form of implicit feedback, where users' watches tell us which things they prefer to see and which they'd rather not see.\n",
        "2. It can also be seen as expressesing how much the users liked the movies they did watch. This is a form of explicit feedback: given that a user watched a movie, we can tell roughly how much they liked by looking at the rating they have given.\n",
        "\n",
        "In this tutorial, we are focusing on a retrieval system: a model that predicts a set of movies from the catalogue that the user is likely to watch. Often, implicit data is more useful here, and so we are going to treat Movielens as an implicit system. This means that every movie a user watched is a positive example, and every movie they have not seen is an implicit negative example."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sawo1x8kQk9b"
      },
      "source": [
        "## Imports\n",
        "\n",
        "\n",
        "Let's first get our imports out of the way."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0vJOdh9WbTpd"
      },
      "outputs": [],
      "source": [
        "!pip install -q tensorflow-recommenders\n",
        "!pip install -q --upgrade tensorflow-datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SZGYDaF-m5wZ"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import pprint\n",
        "import tempfile\n",
        "\n",
        "from typing import Dict, Text\n",
        "\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BxQ_hy7xPH3N"
      },
      "outputs": [],
      "source": [
        "import tensorflow_recommenders as tfrs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5PAqjR4a1RR4"
      },
      "source": [
        "## Preparing the dataset\n",
        "\n",
        "Let's first have a look at the data.\n",
        "\n",
        "We use the MovieLens dataset from [Tensorflow Datasets](https://www.tensorflow.org/datasets). Loading `movie_lens/100k_ratings` yields a `tf.data.Dataset` object containing the ratings data and loading `movie_lens/100k_movies` yields a `tf.data.Dataset` object containing only the movies data.\n",
        "\n",
        "Note that since the MovieLens dataset does not have predefined splits, all data are under `train` split."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aaQhqcLGP0jL"
      },
      "outputs": [],
      "source": [
        "# Ratings data.\n",
        "ratings = tfds.load(\"movie_lens/100k-ratings\", split=\"train\")\n",
        "# Features of all the available movies.\n",
        "movies = tfds.load(\"movie_lens/100k-movies\", split=\"train\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JRHorm8W1yf3"
      },
      "source": [
        "The ratings dataset returns a dictionary of movie id, user id, the assigned rating, timestamp, movie information, and user information:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_1-KQV2ynMdh"
      },
      "outputs": [],
      "source": [
        "for x in ratings.take(1).as_numpy_iterator():\n",
        "  pprint.pprint(x)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qGLGCjSt_q96"
      },
      "source": [
        "The movies dataset contains the movie id, movie title, and data on what genres it belongs to. Note that the genres are encoded with integer labels."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kHLsIHhw_x1d"
      },
      "outputs": [],
      "source": [
        "for x in movies.take(1).as_numpy_iterator():\n",
        "  pprint.pprint(x)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YUdT-f4RxMKs"
      },
      "source": [
        "In this example, we're going to focus on the ratings data. Other tutorials explore how to use the movie information data as well to improve the model quality.\n",
        "\n",
        "We keep only the `user_id`, and `movie_title` fields in the dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uhbEvPJqxLec"
      },
      "outputs": [],
      "source": [
        "ratings = ratings.map(lambda x: {\n",
        "    \"movie_title\": x[\"movie_title\"],\n",
        "    \"user_id\": x[\"user_id\"],\n",
        "})\n",
        "movies = movies.map(lambda x: x[\"movie_title\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Iu4XSa_G1nyN"
      },
      "source": [
        "To fit and evaluate the model, we need to split it into a training and evaluation set. In an industrial recommender system, this would most likely be done by time: the data up to time $T$ would be used to predict interactions after $T$.\n",
        "\n",
        "\n",
        "In this simple example, however, let's use a random split, putting 80% of the ratings in the train set, and 20% in the test set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rS0eDfkjnjJL"
      },
      "outputs": [],
      "source": [
        "tf.random.set_seed(42)\n",
        "shuffled = ratings.shuffle(100_000, seed=42, reshuffle_each_iteration=False)\n",
        "\n",
        "train = shuffled.take(80_000)\n",
        "test = shuffled.skip(80_000).take(20_000)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gVi1HJfR9D7H"
      },
      "source": [
        "Let's also figure out unique user ids and movie titles present in the data. \n",
        "\n",
        "This is important because we need to be able to map the raw values of our categorical features to embedding vectors in our models. To do that, we need a vocabulary that maps a raw feature value to an integer in a contiguous range: this allows us to look up the corresponding embeddings in our embedding tables."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MKROCiPo_5LJ"
      },
      "outputs": [],
      "source": [
        "movie_titles = movies.batch(1_000)\n",
        "user_ids = ratings.batch(1_000_000).map(lambda x: x[\"user_id\"])\n",
        "\n",
        "unique_movie_titles = np.unique(np.concatenate(list(movie_titles)))\n",
        "unique_user_ids = np.unique(np.concatenate(list(user_ids)))\n",
        "\n",
        "unique_movie_titles[:10]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eCi-seR86qqa"
      },
      "source": [
        "## Implementing a model\n",
        "\n",
        "Choosing the architecure of our model a key part of modelling.\n",
        "\n",
        "Because we are building a two-tower retrieval model, we can build each tower separately and then combine them in the final model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z20PyfSXP3Um"
      },
      "source": [
        "### The query tower\n",
        "\n",
        "Let's start with the query tower.\n",
        "\n",
        "The first step is to decide on the dimensionality of the query and candidate representations:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QbIy1FP8aCTq"
      },
      "outputs": [],
      "source": [
        "embedding_dimension = 32"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IJYwjpLRaEzj"
      },
      "source": [
        "Higher values will correspond to models that may be more accurate, but will also be slower to fit and more prone to overfitting.\n",
        "\n",
        "The second is to define the model itself. Here, we're going to use Keras preprocessing layers to first convert user ids to integers, and then convert those to user embeddings via an `Embedding` layer. Note that we use the list of unique user ids we computed earlier as a vocabulary:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kHQZJEhXP93N"
      },
      "outputs": [],
      "source": [
        "user_model = tf.keras.Sequential([\n",
        "  tf.keras.layers.experimental.preprocessing.StringLookup(\n",
        "      vocabulary=unique_user_ids, mask_token=None),\n",
        "  # We add an additional embedding to account for unknown tokens.\n",
        "  tf.keras.layers.Embedding(len(unique_user_ids) + 1, embedding_dimension)\n",
        "])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qvo2pEcdaiec"
      },
      "source": [
        "A simple model like this corresponds exactly to a classic [matrix factorization](https://ieeexplore.ieee.org/abstract/document/4781121) approach. While defining a subclass of `tf.keras.Model` for this simple model might be overkill, we can easily extend it to an arbitrarily complex model using standard Keras components, as long as we return an `embedding_dimension`-wide output at the end."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dG4YFy9SQ08d"
      },
      "source": [
        "### The candidate tower\n",
        "\n",
        "We can do the same with the candidate tower."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qNUwfIJTQ332"
      },
      "outputs": [],
      "source": [
        "movie_model = tf.keras.Sequential([\n",
        "  tf.keras.layers.experimental.preprocessing.StringLookup(\n",
        "      vocabulary=unique_movie_titles, mask_token=None),\n",
        "  tf.keras.layers.Embedding(len(unique_movie_titles) + 1, embedding_dimension)\n",
        "])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r10RiPtqVIAl"
      },
      "source": [
        "### Metrics\n",
        "\n",
        "In our training data we have positive (user, movie) pairs. To figure out how good our model is, we need to compare the affinity score that the model calculates for this pair to the scores of all the other possible candidates: if the score for the positive pair is higher than for all other candidates, our model is highly accurate.\n",
        "\n",
        "To do this, we can use the `tfrs.metrics.FactorizedTopK` metric. The metric has one required argument: the dataset of candidates that are used as implicit negatives for evaluation.\n",
        "\n",
        "In our case, that's the `movies` dataset, converted into embeddings via our movie model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1dLDL6pZVPO8"
      },
      "outputs": [],
      "source": [
        "metrics = tfrs.metrics.FactorizedTopK(\n",
        "  candidates=movies.batch(128).map(movie_model)\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nCaCqJsXSkCo"
      },
      "source": [
        "### Loss\n",
        "\n",
        "The next component is the loss used to train our model. TFRS has several loss layers and tasks to make this easy.\n",
        "\n",
        "In this instance, we'll make use of the `Retrieval` task object: a convenience wrapper that bundles together the loss function and metric computation:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tJ61Iz2QTBw3"
      },
      "outputs": [],
      "source": [
        "task = tfrs.tasks.Retrieval(\n",
        "  metrics=metrics\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9-3xFC-1cbz0"
      },
      "source": [
        "The task itself is a Keras layer that takes the query and candidate embeddings as arguments, and returns the computed loss: we'll use that to implement the model's training loop."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FZUFeSlWRHGx"
      },
      "source": [
        "### The full model\n",
        "\n",
        "We can now put it all together into a model. TFRS exposes a base model class (`tfrs.models.Model`) which streamlines bulding models: all we need to do is to set up the components in the `__init__` method, and implement the `compute_loss` method, taking in the raw features and returning a loss value.\n",
        "\n",
        "The base model will then take care of creating the appropriate training loop to fit our model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8n7c5CHFp0ow"
      },
      "outputs": [],
      "source": [
        "class MovielensModel(tfrs.Model):\n",
        "\n",
        "  def __init__(self, user_model, movie_model):\n",
        "    super().__init__()\n",
        "    self.movie_model: tf.keras.Model = movie_model\n",
        "    self.user_model: tf.keras.Model = user_model\n",
        "    self.task: tf.keras.layers.Layer = task\n",
        "\n",
        "  def compute_loss(self, features: Dict[Text, tf.Tensor], training=False) -\u003e tf.Tensor:\n",
        "    # We pick out the user features and pass them into the user model.\n",
        "    user_embeddings = self.user_model(features[\"user_id\"])\n",
        "    # And pick out the movie features and pass them into the movie model,\n",
        "    # getting embeddings back.\n",
        "    positive_movie_embeddings = self.movie_model(features[\"movie_title\"])\n",
        "\n",
        "    # The task computes the loss and the metrics.\n",
        "    return self.task(user_embeddings, positive_movie_embeddings)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I7B8PdfNqyuN"
      },
      "source": [
        "The `tfrs.Model` base class is a simply convenience class: it allows us to compute both training and test losses using the same method.\n",
        "\n",
        "Under the hood, it's still a plain Keras model. You could achieve the same functionality by inheriting from `tf.keras.Model` and overriding the `train_step` and `test_step` functions (see [the guide](https://keras.io/guides/customizing_what_happens_in_fit/) for details):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z3QywMjqsH4F"
      },
      "outputs": [],
      "source": [
        "class NoBaseClassMovielensModel(tf.keras.Model):\n",
        "\n",
        "  def __init__(self, user_model, movie_model):\n",
        "    super().__init__()\n",
        "    self.movie_model: tf.keras.Model = movie_model\n",
        "    self.user_model: tf.keras.Model = user_model\n",
        "    self.task: tf.keras.layers.Layer = task\n",
        "\n",
        "  def train_step(self, features: Dict[Text, tf.Tensor]) -\u003e tf.Tensor:\n",
        "\n",
        "    # Set up a gradient tape to record gradients.\n",
        "    with tf.GradientTape() as tape:\n",
        "\n",
        "      # Loss computation.\n",
        "      user_embeddings = self.user_model(features[\"user_id\"])\n",
        "      positive_movie_embeddings = self.movie_model(features[\"movie_title\"])\n",
        "      loss = self.task(user_embeddings, positive_movie_embeddings)\n",
        "\n",
        "      # Handle regularization losses as well.\n",
        "      regularization_loss = sum(self.losses)\n",
        "\n",
        "      total_loss = loss + regularization_loss\n",
        "\n",
        "    gradients = tape.gradient(total_loss, self.trainable_variables)\n",
        "    self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))\n",
        "\n",
        "    metrics = {metric.name: metric.result() for metric in self.metrics}\n",
        "    metrics[\"loss\"] = loss\n",
        "    metrics[\"regularization_loss\"] = regularization_loss\n",
        "    metrics[\"total_loss\"] = total_loss\n",
        "\n",
        "    return metrics\n",
        "\n",
        "  def test_step(self, features: Dict[Text, tf.Tensor]) -\u003e tf.Tensor:\n",
        "\n",
        "    # Loss computation.\n",
        "    user_embeddings = self.user_model(features[\"user_id\"])\n",
        "    positive_movie_embeddings = self.movie_model(features[\"movie_title\"])\n",
        "    loss = self.task(user_embeddings, positive_movie_embeddings)\n",
        "\n",
        "    # Handle regularization losses as well.\n",
        "    regularization_loss = sum(self.losses)\n",
        "\n",
        "    total_loss = loss + regularization_loss\n",
        "\n",
        "    metrics = {metric.name: metric.result() for metric in self.metrics}\n",
        "    metrics[\"loss\"] = loss\n",
        "    metrics[\"regularization_loss\"] = regularization_loss\n",
        "    metrics[\"total_loss\"] = total_loss\n",
        "\n",
        "    return metrics"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uHnl2nHMs_E0"
      },
      "source": [
        "In these tutorials, however, we stick to using the `tfrs.Model` base class to keep our focus on modelling and abstract away some of the boilerplate."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yDN_LJGlnRGo"
      },
      "source": [
        "## Fitting and evaluating\n",
        "\n",
        "After defining the model, we can use standard Keras fitting and evaluation routines to fit and evaluate the model.\n",
        "\n",
        "Let's first instantiate the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aW63YaqP2wCf"
      },
      "outputs": [],
      "source": [
        "model = MovielensModel(user_model, movie_model)\n",
        "model.compile(optimizer=tf.keras.optimizers.Adagrad(learning_rate=0.1))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nma0vc2XdN5g"
      },
      "source": [
        "Then shuffle, batch, and cache the training and evaluation data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "53QJwY1gUnfv"
      },
      "outputs": [],
      "source": [
        "cached_train = train.shuffle(100_000).batch(8192).cache()\n",
        "cached_test = test.batch(4096).cache()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u8mHTxKAdTJO"
      },
      "source": [
        "Then train the  model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZxPntlT8EFOZ"
      },
      "outputs": [],
      "source": [
        "model.fit(cached_train, epochs=3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YsluR8audV9W"
      },
      "source": [
        "As the model trains, the loss is falling and a set of top-k retrieval metrics is updated. These tell us whether the true positive is in the top-k retrieved items from the entire candidate set. For example, a top-5 categorical accuracy metric of 0.2 would tell us that, on average, the true positive is in the top 5 retrieved items 20% of the time.\n",
        "\n",
        "Note that, in this example, we evaluate the metrics during training as well as evaluation. Because this can be quite slow with large candidate sets, it may be prudent to turn metric calculation off in training, and only run it in evaluation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7Gxp5RLFcv64"
      },
      "source": [
        "Finally, we can evaluate our model on the test set:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W-zu6HLODNeI"
      },
      "outputs": [],
      "source": [
        "model.evaluate(cached_test, return_dict=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JKZyP9A1dxit"
      },
      "source": [
        "Test set performance is much worse than training performance. This is due to two factors:\n",
        "\n",
        "1. Our model is likely to perform better on the data that it has seen, simply because it can memorize it. This overfitting phenomenon is especially strong when models have many parameters. It can be mediated by model regularization and use of user and movie features that help the model generalize better to unseen data.\n",
        "2. The model is re-recommending some of users' already watched movies. These known-positive watches can crowd out test movies out of top K recommendations.\n",
        "\n",
        "The second phenomenon can be tackled by excluding previously seen movies from test recommendations. This approach is relatively common in the recommender systems literature, but we don't follow it in these tutorials. If not recommending past watches is important, we should expect appropriately specified models to learn this behaviour automatically from past user history and contextual information. Additionally, it is often appropriate to recommend the same item multiple times (say, an evergreen TV series or a regularly purchased item)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NB2v43NJU3Xf"
      },
      "source": [
        "## Making predictions\n",
        "\n",
        "Now that we have a model, we would like to be able to make predictions. We can use the `tfrs.layers.factorized_top_k.BruteForce` layer to do this."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IRD6bEtZW_8j"
      },
      "outputs": [],
      "source": [
        "# Create a model that takes in raw query features, and\n",
        "index = tfrs.layers.factorized_top_k.BruteForce(model.user_model)\n",
        "# recommends movies out of the entire movies dataset.\n",
        "index.index(movies.batch(100).map(model.movie_model), movies)\n",
        "\n",
        "# Get recommendations.\n",
        "_, titles = index(tf.constant([\"42\"]))\n",
        "print(f\"Recommendations for user 42: {titles[0, :3]}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cOoLv6ZMKg0L"
      },
      "source": [
        "Of course, the `BruteForce` layer is going to be too slow to serve a model with many possible candidates. The following sections shows how to speed this up by using an approximate retrieval index."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pFlvp5DK4Ow8"
      },
      "source": [
        "## Model serving\n",
        "\n",
        "After the model is trained, we need a way to deploy it.\n",
        "\n",
        "In a two-tower retrieval model, serving has two components:\n",
        "\n",
        "- a serving query model, taking in features of the query and transforming them into a query embedding, and\n",
        "- a serving candidate model. This most often takes the form of an approximate nearest neighbours (ANN) index which allows fast approximate lookup of candidates in response to a query produced by the query model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NmhfltdpWZ06"
      },
      "source": [
        "### Exporting the entire model\n",
        "\n",
        "In TFRS, both components can be packaged into a single exportable model, giving us a model that takes the raw user id and returns the titles of top movies for that user. This is done via exporting the model to a `SavedModel` format, which makes it possible to serve using [TensorFlow Serving](https://www.tensorflow.org/tfx/guide/serving).\n",
        "\n",
        "To deploy a model like this, we simply export the `BruteForce` layer we created above:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oJkRNBfCW5_E"
      },
      "outputs": [],
      "source": [
        "# Export the query model.\n",
        "with tempfile.TemporaryDirectory() as tmp:\n",
        "  path = os.path.join(tmp, \"model\")\n",
        "\n",
        "  # Save the index.\n",
        "  index.save(path)\n",
        "\n",
        "  # Load it back; can also be done in TensorFlow Serving.\n",
        "  loaded = tf.keras.models.load_model(path)\n",
        "\n",
        "  # Pass a user id in, get top predicted movie titles back.\n",
        "  scores, titles = loaded([\"42\"])\n",
        "\n",
        "  print(f\"Recommendations: {titles[0][:3]}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PIJ2GK9k-vHe"
      },
      "source": [
        "### Exporting the query model\n",
        "\n",
        "We can also export the query model and the candidates separately. This is useful if you'd like to pair your query model with an external approximate nearest neighbours system.\n",
        "\n",
        "To export the query model, you can do the following:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wi72fBjan0JU"
      },
      "outputs": [],
      "source": [
        "# Export the query model.\n",
        "with tempfile.TemporaryDirectory() as tmp:\n",
        "  path = os.path.join(tmp, \"query_model\")\n",
        "  model.user_model.save(path)\n",
        "  loaded = tf.keras.models.load_model(path)\n",
        "  query_embedding = loaded(tf.constant([\"10\"]))\n",
        "\n",
        "  print(f\"Query embedding: {query_embedding[0, :3]}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O-ENamtGmLHY"
      },
      "source": [
        "### Building a candidate ANN index\n",
        "\n",
        "Exporting candidate representations is more involved. Firstly, we want to pre-compute them to make sure serving is fast; this is especially important if the candidate model is computationally intensive (for example, if it has many or wide layers; or uses complex representations for text or images). Secondly, we would like to take the precomputed representations and use them to construct a fast approximate retrieval index.\n",
        "\n",
        "\n",
        "We can use [Annoy](https://github.com/spotify/annoy) to build such an index.\n",
        "\n",
        "Annoy isn't included in the base TFRS package. To install it, run:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "geCJgN9YOcUD"
      },
      "outputs": [],
      "source": [
        "!pip install -q annoy"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y8AAuhapJZ_j"
      },
      "source": [
        "We can now create the index object."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jghIJ0ydqT6R"
      },
      "outputs": [],
      "source": [
        "from annoy import AnnoyIndex\n",
        "\n",
        "index = AnnoyIndex(embedding_dimension, \"dot\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YdojJq9m5kmZ"
      },
      "source": [
        "Then take the candidate dataset and transform its raw features into embeddings using the movie model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gZnsvraK5v_a"
      },
      "outputs": [],
      "source": [
        "movie_embeddings = movies.enumerate().map(lambda idx, title: (idx, title, model.movie_model(title)))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AStmvt_I5_hx"
      },
      "source": [
        "And then index the movie_id, movie embedding pairs into our Annoy index:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H9olnuVrmZxa"
      },
      "outputs": [],
      "source": [
        "movie_id_to_title = dict((idx, title) for idx, title, _ in movie_embeddings.as_numpy_iterator())\n",
        "\n",
        "# We unbatch the dataset because Annoy accepts only scalar (id, embedding) pairs.\n",
        "for movie_id, _, movie_embedding in movie_embeddings.as_numpy_iterator():\n",
        "  index.add_item(movie_id, movie_embedding)\n",
        "\n",
        "# Build a 10-tree ANN index.\n",
        "index.build(10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y04_XP_j7CIn"
      },
      "source": [
        "We can then retrieve nearest neighbours:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4U3vr8aWnPTm"
      },
      "outputs": [],
      "source": [
        "for row in test.batch(1).take(3):\n",
        "  query_embedding = model.user_model(row[\"user_id\"])[0]\n",
        "  candidates = index.get_nns_by_vector(query_embedding, 3)\n",
        "  print(f\"Candidates: {[movie_id_to_title[x] for x in candidates]}.\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ws4w5jQ3fX87"
      },
      "source": [
        "## Next steps\n",
        "\n",
        "This concludes the retrieval tutorial.\n",
        "\n",
        "To expand on what is presented here, have a look at:\n",
        "\n",
        "1. Learning multi-task models: jointly optimizing for ratings and clicks.\n",
        "2. Using movie metadata: building a more complex movie model to alleviate cold-start."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i0UFyyocksOF"
      },
      "outputs": [],
      "source": [
        ""
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "basic_retrieval.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
